"""
Created on 10-Dec-2009

@author: dexen
"""

import sys
import time
import multiprocessing as mp

import rpyc


from dexen.system.common.base import Action
from dexen.system.common.meta import JobMeta
from dexen.system.common.service import DexenAddr
from dexen.system.node.master.env import MasterEnv
from dexen.system.node.master.proxy import SlaveProxy
from dexen.system.common.transact import (SlaveSendTransact, SlaveRecvTransact, 
                                          Transaction, MasterRecvTransact)
from dexen.system.node.master.mgrs.state_mgr import StateMgr
from dexen.system.node.master.executor import MasterTaskExecutor


#===============================================================================
# Globals
#===============================================================================
logger = None
        
#===============================================================================
# Slave Manager
#===============================================================================
class MasterCore(mp.Process):
    def __init__(self, job_meta, comm, reg_addr):
        assert isinstance(job_meta, JobMeta)        
        assert isinstance(reg_addr, DexenAddr)
        mp.Process.__init__(self)
        self.job_meta = job_meta
        self.comm = comm
        self.reg_addr = reg_addr # resource registry address
        self.slaves = {} # slave proxies
        self.scheduled = set() # names of the currently scheduled slaves
        self.state_mgr = StateMgr()
        self.act_id = 1
        self.state_mgr.set_cur_time(self.act_id)
        self.master_env = MasterEnv(self.state_mgr)
        self.m_task_executor = MasterTaskExecutor(job_meta, self.master_env)
                           
    def _get_slave_addr(self, slave_name):
        reg_addr = self.reg_addr
        assert isinstance(reg_addr, DexenAddr)
        conn = reg_addr.get_conn()
        addr = conn.root.get_slave_addr(slave_name)
        addr = rpyc.classic.obtain(addr)
        conn.close()
        return addr
    
    def _schedule(self, slave_name, transact, master_state):
        assert isinstance(transact, SlaveSendTransact)
        self.scheduled.add(slave_name)
        if not self.slaves.has_key(slave_name): 
            slave_addr = self._get_slave_addr(slave_name)
            print "I am creting SlaveProxy for ", slave_name
            self.slaves[slave_name] = SlaveProxy(slave_name, slave_addr, self.job_meta)
        slave = self.slaves[slave_name]
        assert isinstance(slave, SlaveProxy)       
        slave.execute(transact, master_state)
        self._create_action(transact)

    def _create_action(self, transact):
        master_state = self.state_mgr.get_master_state()
        assert isinstance(transact, Transaction)
        transact.strip_inds()
        action = Action(self.act_id, master_state, transact)
        self.act_id += 1
        self.state_mgr.set_cur_time(self.act_id)
        self.comm.add_action(action)

    # TODO: cater for crashes
    def _update_slaves(self):
        #print "updating slaves"
        for slave_name in self.scheduled.copy():
            slave = self.slaves[slave_name]
            assert isinstance(slave, SlaveProxy)

            slave.update()

            if slave.is_busy():
                # This means that after querying the slave, the slave hasn't 
                # finished executing the task.
                continue

            self.scheduled.remove(slave_name)
            
            transact = slave.get_transact()
            sent_transact = slave.get_sent_transact()
            assert isinstance(transact, SlaveRecvTransact)
            assert isinstance(sent_transact, SlaveSendTransact)
            

            
            if slave.is_crashed():
                print slave.name, "crashed"
                self.state_mgr.recover_s_transact(transact, sent_transact)
                self._create_action(transact)
            if slave.is_idle():
                """
                If the slave is in idle state, it means that the proxy has just
                got back the results pending to be processed
                """
                print slave.name, "has become idle", time.time()
                self.state_mgr.set_s_transact(transact, sent_transact)

                self._create_action(transact)

    def _update_master(self):
        transact = self.m_task_executor.execute()
        if transact:
            assert isinstance(transact, MasterRecvTransact)
            self.state_mgr.set_m_transact(transact)
            self._create_action(transact)
            if transact.task_failed():
                self.state_mgr.stop()
          
    def run(self):       
        self.m_task_executor.setup()     
        sys.path.append(self.job_meta.root_dir)
        while True:
            if self.comm.has_stopped():
                break
            self._update_slaves()
            if self.comm.has_slave():
                print "MasterCore: Has slave"
                transact = self.state_mgr.get_s_transact()
                if transact != None:
                    master_state = self.state_mgr.get_master_state()
                    slave_name = self.comm.get_slave()
                    self._schedule(slave_name, transact, master_state)    
            self._update_master()
            if self.state_mgr.has_finished():
                break
            time.sleep(1)
                
